from torch.utils.data import Dataset, DataLoader
import torch

class InteractionDataset(Dataset):
    def __init__(self, user_ids, item_ids, labels):
        self.user_ids = user_ids
        self.item_ids = item_ids
        self.labels = labels
        
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        return {
            "user_id": self.user_ids[idx],
            "item_id": self.item_ids[idx],
            "label": self.labels[idx]
        }

def create_data_loaders(config):
    from data.generate_data import DataGenerator
    
    generator = DataGenerator()
    user_ids, item_ids, labels = generator.generate_interactions()
    
    # 分割数据集
    dataset_size = len(user_ids)
    test_size = int(config.TEST_SIZE * dataset_size)
    train_size = dataset_size - test_size
    
    # 创建数据集
    train_dataset = InteractionDataset(
        user_ids[:train_size],
        item_ids[:train_size],
        labels[:train_size]
    )
    
    test_dataset = InteractionDataset(
        user_ids[train_size:],
        item_ids[train_size:],
        labels[train_size:]
    )
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset, 
        batch_size=config.BATCH_SIZE, 
        shuffle=True,
        num_workers=4,
        pin_memory=torch.cuda.is_available()
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=config.BATCH_SIZE * 2,
        shuffle=False,
        num_workers=2,
        pin_memory=torch.cuda.is_available()
    )
    
    return train_loader, test_loader